import random
import numpy as np
from scipy.stats import ortho_group

from view.endSrc.MixedGaussianDataset import MixedGaussianDataset
from view.endSrc.GaussianClusters import GaussianClusters

class MixedGaussianDatasetGenerator():
    '''
    This class provide methods to generate samples from multivariate normal distribution with 
    various configurations for the diversification of datasets.
    The objective is to generate datasets that range from simple to complex by increasing the 
    number of instances, the number of features, and the number of clusters and setting 
    various overlapping thresholds and the range of the mean values.
    '''


    def __init__(self, numOfInstances, numOfFeatures, overlap, paiList, xRange = None):
        
        '''
        Initialize all member variables.

        Parameters
        --------
        self.N: number of instances
        self.d: number of features
        self.k: number of clusters
        self.v overlapping thresholds
        self.clusterLabels: a list containing the cluster labels index
        if k = 3, self.clusterLabels may like [0, 0, 0, 1, 1, 1, 2, 2, 2]
        self.dataset: a N*d dataset array
        self.clusterCen: a list containing centers of different clusters
        self.numOfInstancesEach: number of instances for each components
        self.xRange: a list represents ranges for center values.

        Examples:
        --------
        datasetGen = MixedGaussianDatasetGenerator()
        '''

        for p in paiList:
            assert p > 0, 'every element in paiList must be positive'

        self.n = numOfInstances
        self.d = numOfFeatures
        self.k = len(paiList)
        self.overlap = overlap

        self.meanList = None
        self.varMatList = None
        self.labels = None          # store generated labels from recently called self.genMixedDataset
        self.mgDataset = None         # MixedGaussianDataset

        # Normalize the pai value to ensure the summation of all pais equals to one
        self.paiList = [pai/sum(paiList) for pai in paiList]
        self.numOfInstancesEach = [int(pai * self.n) for pai in self.paiList]
        assert all(n > 0 for n in self.numOfInstancesEach), 'numOfInstances is too small, please try again'


        # Ensure the sum of instances of all components is equal to the total number of instances
        if sum(self.numOfInstancesEach) != self.n:
            self.numOfInstancesEach[0] += self.n - sum(self.numOfInstancesEach)

        # In the original alogrithm, the xRange is set as [1, self.k] compulsively which can not be changed.
        # In this case, user can set the xRange, and the default value of xRange is [1, self.k].
        self.xRange = xRange
        if xRange == None:
            self.xRange = [1, self.k]

    def genSeparableComponents(self):

        '''
        This method generate the mean and covariance given certain parameter settings for one cluster,
        which are stored in self.meanList and self.varMatList.
        

        Examples:
        --------
        datasetGen.genSeparableComponents()
        '''

        xMin = self.xRange[0]
        xMax = self.xRange[1]
        self.meanList = []

        while len(self.meanList) < self.k:
            vec = []
            for dimIndex in range(self.d):
                vec.append(random.randint(xMin, xMax))
                # gaussianVector.append(random.uniform(xMin, xMax))

            # Ensure that the each vector(only list works, not numpy) in gaussian components is unique
            if vec not in self.meanList:
                self.meanList.append(vec)

        # transform to list of numpy
        for i, v in enumerate(self.meanList):
            self.meanList[i] = np.asarray(v) - 0.5 # self.clusterCen is of size k*d


        # Compare the distance between cluster centers and calculate the projection of distances
        # on all dimensions to meet the requirement sigma_i > overlap * min_j{distance_proj_on_axis=i(a,j)}
        radiusList = []
        for clusterInd in range(self.k):
            clusterDistList = [] # This contains the euclidean difference
            for anotherClusterInd in range(self.k):
                if clusterInd != anotherClusterInd:
                    dist = self.meanList[clusterInd] - self.meanList[anotherClusterInd] # dist is 1 * d
                    clusterDistList.append(dist)
            radius = []
            for i in range(self.d):
                v = np.zeros(self.d)
                v[i] = 1  # norm of v is 1
                projList = []
                for j in range(self.k-1):
                    proj = (np.dot(clusterDistList[j], v)) * v   # dot product with e_i = [0, 0, 1_i, 0, 0]
                    projNorm = np.linalg.norm(proj)
                    projList.append(projNorm)
                radius.append(min(projList))
            radiusList.append(radius)

        self.varMatList = []
        for ind in range(self.k):
            varList = [(rad * self.overlap) for rad in radiusList[ind]]

            mat= np.diag(varList)   # first create a positive diagonal matrix

            A = np.float32(ortho_group.rvs(dim=self.d)) # A is an orthogonal matrix
            Amat = np.dot(A, mat)
            mat = np.dot(Amat, A.transpose()) # Ensure covariance matrix is positive definite

            self.varMatList.append(mat)


    def genMixedDataset(self):
        
        '''
        This method generate the dataset for different clusters given the mean and covariance 
        for different clusters and the ratio(pai) for different clusters, then merge them into
        a single dataset with N instances. In this way, a dataset with N instances, d features,
        and k clusters is generated. Also, a relationship array for gaussian clusters is generated
        to record the true cluster labels of dataset.


        Examples:
        --------
        dataset = datasetGen.genMixedDataset()
        '''

        # firstly, generate meanList, varMatList for Gaussian distribution
        if self.meanList is None:
            self.genSeparableComponents()


        print('numOfInstances: ', self.numOfInstancesEach)

        # generate observations from meanList, varMatList
        dataset = []
        self.labels = []
        for ind in range(self.k):
            mean = self.meanList[ind]
            cov = self.varMatList[ind]
            size = self.numOfInstancesEach[ind]

            data = np.random.multivariate_normal(mean, cov, size)
            dataset.append(data)
            self.labels += ([ind] * size)
        dataset = np.concatenate(dataset, axis = 0)

        print('self.k: ', self.k)
        print('self.labels: ', self.labels)

        # build output object
        gc = GaussianClusters(self.meanList, self.varMatList, self.paiList, labels=self.labels)
        self.mgDataset = MixedGaussianDataset(dataset, self.labels, gc)

        return self.mgDataset
    

    # def display(self):
    #
    #     '''
    #     This method provide ways to print out dataset, relationship array for gaussian clusters,
    #     and the shape of above two arrays. It also provides ways to print out the centers(means)
    #     and the covariance matrix of the k clusters.
    #
    #
    #     Examples:
    #     --------
    #     dataset.display()
    #     '''
    #
    #     print('dataset: ', self.mgDataset)
    #     print('Gaussian labels: ', self.labels)
    #     print('Cluster Centers: ', self.meanList)
    #
    #     print('Covariance matrixes: ')
    #     for mt in self.varMatList:
    #         print(mt)

            
# if __name__ == '__main__':
#
#     ps = []
#     for i in range(10):
#         ps.append(random.uniform(1, 3))
#
#     print(ps)
#     g = MixedGaussianDatasetGenerator(numOfInstances=2000, numOfFeatures=2,
#                                       overlap=1.6,
#                                       paiList=ps,
#                                       xRange=[10, 100])
#
#     g.genSeparableComponents()
#     mg = g.genMixedDataset()
#
#     g.display()
#
#     g.visualization()
#     g.visual2D()

